#include "PlaneMesh.h"

using namespace OBR;

std::shared_ptr<CudaProgram> PlaneMesh::program;
CudaBuffer PlaneMesh::sbt_data_payload;

PlaneMesh::PlaneMesh() : Mesh("shaders/Mesh/PlaneMesh.cu")
{
    // init program opt & desc
    this->program_desc.hitgroup.entryFunctionNameCH = "__closesthit__plane_mesh";
    this->program_desc.hitgroup.entryFunctionNameIS = "__intersection__plane_mesh";
    this->program_desc.kind = OPTIX_PROGRAM_GROUP_KIND_HITGROUP;
    this->program_desc.flags = OPTIX_GEOMETRY_FLAG_NONE;
}

PlaneMesh::~PlaneMesh()
{
}

void PlaneMesh::_buildFromConfig(const Json::Value &config)
{
    this->U = ConfigHelpers::loadFloat3(config["U"]);
    this->V = ConfigHelpers::loadFloat3(config["V"]);
    this->O = ConfigHelpers::loadFloat3(config["O"]);
}

OptixBuildInput PlaneMesh::getASBuildInput()
{
    OptixBuildInput build_input = {};

    // build flags
    if (build_flags)
    {
        delete[] build_flags;
    }
    build_flags = new uint32_t[this->ref_cnt];
    for (int i = 0; i < this->ref_cnt; i++)
    {
        // TODO: specify build flags
        build_flags[i] = OPTIX_GEOMETRY_FLAG_DISABLE_ANYHIT;
    }

    // set build input
    build_input.type = OPTIX_BUILD_INPUT_TYPE_CUSTOM_PRIMITIVES;
    build_input.customPrimitiveArray.flags = build_flags;
    build_input.customPrimitiveArray.numPrimitives = 1;
    build_input.customPrimitiveArray.numSbtRecords = 1;
    build_input.customPrimitiveArray.aabbBuffers =
        reinterpret_cast<CUdeviceptr *>(&(this->d_aabb.d_ptr));

    // index offset
    uint sbt_idx_offset = 0;
    this->d_sbt_idx_offsets.resize(sizeof(uint));
    this->d_sbt_idx_offsets.upload(&sbt_idx_offset, sizeof(uint));
    build_input.customPrimitiveArray.sbtIndexOffsetSizeInBytes = sizeof(uint);
    build_input.customPrimitiveArray.sbtIndexOffsetBuffer = this->d_sbt_idx_offsets.d_pointer();

    return build_input;
}

OptixAccelBuildOptions PlaneMesh::getASBuildOptions()
{
    OptixAccelBuildOptions build_options = {};

    build_options.buildFlags = OPTIX_BUILD_FLAG_ALLOW_COMPACTION;
    build_options.operation = OPTIX_BUILD_OPERATION_BUILD;

    return build_options;
}

void PlaneMesh::uploadVirtualProgramPayload(
    std::vector<std::shared_ptr<VirtualProgramEntry>> instances)
{

    const size_t n_instances = instances.size();
    const size_t data_size = sizeof(PlaneMeshData);
    const size_t memory_size = n_instances * data_size;

    // collect data
    PlaneMeshData *tmpdata = new PlaneMeshData[n_instances];
    for (size_t i = 0; i < n_instances; i++)
    {
        std::shared_ptr<PlaneMesh> m = std::dynamic_pointer_cast<PlaneMesh>(instances[i]);
        tmpdata[i].n = this->N;
        tmpdata[i].u = this->U;
        tmpdata[i].v = this->V;
        tmpdata[i].o = this->O;
    }

    // upload data and free tmp
    PlaneMesh::sbt_data_payload.resize(memory_size);
    PlaneMesh::sbt_data_payload.upload(tmpdata, memory_size);
    delete[] tmpdata;
}

CUdeviceptr PlaneMesh::getPayloadDevicePtr()
{
    return PlaneMesh::sbt_data_payload.d_pointer();
}

void PlaneMesh::setProgram(std::shared_ptr<CudaProgram> prog)
{
    PlaneMesh::program = prog;
}

std::shared_ptr<CudaProgram> PlaneMesh::getProgram()
{
    return PlaneMesh::program;
}

OptixProgramGroupOptions &PlaneMesh::getProgramGroupOptions()
{
    return this->program_opt;
}

OptixProgramGroupDesc &PlaneMesh::getProgramGroupDesc()
{
    this->program_desc.hitgroup.moduleCH = this->cuda_module->optix_module;
    this->program_desc.hitgroup.moduleIS = this->cuda_module->optix_module;
    return this->program_desc;
}

void PlaneMesh::_onBuildEnd()
{
    /* -------------------------- compute N & area -------------------------- */
    float3 UcrossV = cross(this->U, this->V);
    this->N = normalize(UcrossV);
    this->area = length(UcrossV);

    /* ----------------------- compute and upload AABB ---------------------- */
    this->updateAABB();
}

void PlaneMesh::updateAABB()
{
    OptixAabb p;
    this->aabb = Bounds::fromPoint(this->O);
    p = Bounds::fromPoint(this->O + this->U);
    this->aabb = Bounds::combine(this->aabb, p);
    p = Bounds::fromPoint(this->O + this->V);
    this->aabb = Bounds::combine(this->aabb, p);
    p = Bounds::fromPoint(this->O + this->U + this->V);
    this->aabb = Bounds::combine(this->aabb, p);
    this->d_aabb.resize(sizeof(OptixAabb));
    this->d_aabb.upload(&this->aabb, sizeof(OptixAabb));
}